import torch
#transforms 对图像进行处理
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
#使用Relu函数
import torch.nn.functional as F
import torch.optim as optim

# prepare dataset
#设置minibatch大小
batch_size = 64
#compose:组合操作 执行括号里面的操作
#transforms.ToTensor()会把HWC会变成C *H *W（拓展：格式为(h,w,c)，像素顺序为RGB）CWH方便计算 我们这里的MINIST数据集是单通道
#transforms.Normalize 归一化把{0,1,...,255}映射到[0,1]方便计算 第一个参数为均值 第二个参数为方差 因为我们是单通道 则小括号中只用一个参数
"""
简单来说就是将数据按通道进行计算，将每一个通道的数据先计算出其方差与均值，然后再将其每一个通道内的每一个数据减去均值，再除以方差，
得到归一化后的结果。
在深度学习图像处理中，标准化处理之后，可以使数据更好的响应激活函数，提高数据的表现力，减少梯度爆炸和梯度消失的出现。
"""
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
#minist是一个手写数字的训练集 像素28*28即784
train_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
test_dataset = datasets.MNIST(root='../dataset/mnist/', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)


# design model using class


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.l1 = torch.nn.Linear(784, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)

    def forward(self, x):
        #view函数用于改变张量形状 从(n,1,28,28)转为(n,1*28*28)即(n,784) n为分组大小即batch_size
        x = x.view(-1, 784)  # -1其实就是自动获取mini_batch
        #激活函数为relu函数
        """
        torch.nn.Linear的call函数计算公式 output = activation(dot(input, weight) + bias)
        output = activation(dot(input, weight) + bias)
        其中，activation是该层应用的激活函数（默认为ReLU) , dot(input, weight)计算的是输入和权重的点积，bias是偏置项。
        weight和bias的初始值为随机生成的
        """
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        x = F.relu(self.l4(x))
        # 最后一层不做激活，不进行非线性变换 最后一层输入到softmax函数无需激活
        return self.l5(x)


model = Net()

# construct loss and optimizer
#CrossEntropyLoss完成了softmax计算 完成了交叉熵损失
criterion = torch.nn.CrossEntropyLoss()
#SGD随机梯度下降 momentum为冲量 使用momentum可以减缓震荡 加快学习
"""
“冲量”这个概念源自于物理中的力学，表示力对时间的积累效应。
在普通的梯度下降法x+=v
中，每次x的更新量v为v=−dx∗lr，其中dx为目标函数func(x)对x的一阶导数，。
当使用冲量时，则把每次x的更新量v考虑为本次的梯度下降量−dx∗lr与上次x的更新量v乘上一个介于[0,1]的因子momentum的和，即
v ′=−dx∗lr+v∗momemtum
当本次梯度下降- dx * lr的方向与上次更新量v的方向相同时，上次的更新量能够对本次的搜索起到一个正向加速的作用。
当本次梯度下降- dx * lr的方向与上次更新量v的方向相反时，上次的更新量能够对本次的搜索起到一个减速的作用。
"""
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

# training cycle forward, backward, update


def train(epoch):
    running_loss = 0.0
    #enumerate为枚举 第二个参数0为开始下标
    print(type(train_loader))
    for batch_idx, data in enumerate(train_loader, 0):
        # 获得一个批次即mini_batch(这里minibatch=64)的数据和标签
        #将data第一个元素给inputs 第二个元素给target
        inputs, target = data
        #input ([64, 1, 28, 28]) ouput([64]) len(data)=2
        optimizer.zero_grad()
        # 获得模型预测结果(64, 10)
        outputs = model(inputs)
        # 交叉熵代价函数outputs(64,10),target（64）计算预测值和标签的差距
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        #获取交叉熵代价函数的损失值
        running_loss += loss.item()
        #每300次输出 即每计算300*64个数据输出
        if batch_idx % 300 == 299:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0


#test只用正向传播 计算成功了多少个样本
def test():
    correct = 0
    total = 0
    #with torch.no_grad() 即不会计算梯度 我们现在在测设环境，用来评估模型性能，所以要关闭自动更新梯度
    """
    在上述代码中，with torch.no_grad()使用了with语句，这是因为torch. no_grad()是一个上下文管理器。
    这个上下文管理器的主要作用是临时关闭梯度计算，以节省内存并提高计算效率。当代码块执行完毕后，with语句会
    自动结束该上下文，使得梯度计算恢复默认状态。
    """
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            #求每一行最大值的下标 每一行都是一个样本
            """
            在Python中，下划线_通常被用作占位符，表示我们不关心变量的值，或者我们暂时不需要使用该值。在这种情况下，
            _作为torch. ma返回的元组的第一个元素，但是我们不关心它的值，所以我们只使用_来接收这个我们不关心的值。
            """
            _, predicted = torch.max(outputs.data, dim=1)  # dim = 1 行是第0个维度，列是第1个维度
            #label.size输出一个元组（n,1） 取第一个元组中的第一个元素即求样本数n
            total += labels.size(0)
            #predicted == labels真则为1，假则为0，计算最后correct的值
            correct += (predicted == labels).sum().item()  # 张量之间的比较运算
    print('accuracy on test set: %d %% ' % (100 * correct / total))


"""
一个python文件通常有两种使用方法，第一是作为脚本直接执行，第二是 import 到其他的 python 脚本中被调用（模块重用）执行。
因此 if __name__ == 'main': 的作用就是控制这两种情况执行代码的过程，在 if __name__ == 'main': 下的代码只有在第一
种情况下（即文件作为脚本直接执行）才会被执行，而 import 到其他脚本中是不会被执行的
"""
if __name__ == '__main__':
    #训练10轮
    for epoch in range(10):
        train(epoch)
        test()